#include <algorithm>
#include <fstream>
#include <iostream>
#include <crypto/cipher.h>

#include "utils/enums.h"

namespace crypto {
Cipher::Cipher() = default;

bool Cipher::chooseCipher(utils::CipherAlgorithm algorithm, Botan::Cipher_Dir direction) {
	switch (algorithm) {
		case utils::CipherAlgorithm::Kuznyechik_CBC:
			cipher = Botan::Cipher_Mode::create("Kuznyechik/CBC", direction);
		break;
		case utils::CipherAlgorithm::Magma_CBC:
			cipher = Botan::Cipher_Mode::create("GOST-28147-89/CBC", direction);
		break;
		case utils::CipherAlgorithm::AES_256_CBC:
			cipher = Botan::Cipher_Mode::create("AES-256/CBC", direction);
		break;
		default:
			return false; // Unsupported algorithm
	}
	return true;
}

void Cipher::generateIvWithSaving(std::size_t size) {
	Botan::AutoSeeded_RNG rng;
	Botan::InitializationVector iv(rng, size);

	this->iv = iv;
}

Botan::InitializationVector Cipher::generateIV(std::size_t size)
{
	Botan::AutoSeeded_RNG rng;
	Botan::InitializationVector iv(rng, size);
	return iv;
}

void Cipher::setIV(const Botan::secure_vector<uint8_t> &iv) {
	this->iv = Botan::InitializationVector(iv);
}

void Cipher::saveIVinFile(const std::string &filename) const {
	std::ofstream outfile(filename, std::ios::binary);
	if (!outfile.is_open()) {
		throw std::runtime_error("Error: Unable to open file for writing");
	}
	auto save_iv = this->iv.bits_of();
	outfile.write(reinterpret_cast<const char*>(save_iv.data()), static_cast<long>(iv.size()));
}

void Cipher::saveIVInByteArray(Botan::secure_vector<uint8_t> &data) const
{
	auto save_iv = this->iv.bits_of();
	data.insert(data.begin(),save_iv.begin(), save_iv.end());
}

bool Cipher::loadIV(const std::string &filename, std::size_t ivSize) {
	std::ifstream infile(filename, std::ios::binary);
	if (!infile.is_open()) {
		//throw std::runtime_error("Error: Unable to open file for reading");
		std::cerr << "Error: Unable to open file for reading" << std::endl;
		return false;
	}
	std::vector<char, Botan::secure_allocator<char>> saved_iv;
	std::copy_n(std::istreambuf_iterator<char>(infile), ivSize, std::back_inserter(saved_iv));
	this->iv = Botan::InitializationVector(Botan::secure_vector<uint8_t>(saved_iv.begin(), saved_iv.end()));
	return true;
}

Botan::secure_vector<uint8_t> Cipher::encrypt(const Botan::SymmetricKey &key,
		const Botan::secure_vector<uint8_t> &plaintext) const {

	if (!cipher) {
		throw std::runtime_error("No cipher algorithm chosen");
	}
	if (iv.empty()) {
		throw std::runtime_error("IV not set");
	}

	cipher->set_key(key);
	Botan::secure_vector<uint8_t> ciphertext(plaintext.begin(), plaintext.end());
	cipher->start(iv);
	cipher->finish(ciphertext);
	return ciphertext;
}

Botan::secure_vector<uint8_t> Cipher::decrypt(const Botan::SymmetricKey &key,
		const Botan::secure_vector<uint8_t> &ciphertext) const {

	if (!cipher) {
		throw std::runtime_error("No cipher algorithm chosen");
	}
	if (iv.empty()) {
		throw std::runtime_error("IV not set");
	}

	cipher->set_key(key);
	Botan::secure_vector<uint8_t> decrypted(ciphertext.begin(), ciphertext.end());
	cipher->start(iv);
	cipher->finish(decrypted);
	return decrypted;
}

Botan::secure_vector<uint8_t> Cipher::updateEncryption(const Botan::SymmetricKey &key,
													   const Botan::InitializationVector &newIV,
													   const Botan::secure_vector<uint8_t> &plaintext,
													   utils::CipherAlgorithm algorithm) {

	if (!cipher) {
		throw std::runtime_error("No cipher algorithm chosen");
	}
	if (iv.empty()) {
		throw std::runtime_error("old IV not set");
	}
	if (newIV.empty()) {
		throw std::runtime_error("new IV not set");
	}

	chooseCipher(algorithm,  Botan::Cipher_Dir::Decryption);

	cipher->set_key(key);
	Botan::secure_vector<uint8_t> decrypted(plaintext.begin(), plaintext.end());
	cipher->start(iv);
	cipher->finish(decrypted);
	chooseCipher(algorithm,  Botan::Cipher_Dir::Encryption);
	cipher->set_key(key);
	cipher->start(newIV);
	cipher->finish(decrypted);
	return decrypted;
}
}
